"""Registration of Django models with Reversion."""


from django.db import models
from django.db.models.query import QuerySet

from reversion.storage import VersionFileStorageWrapper


class RegistrationError(Exception):
    
    """Exception thrown when registration with Reversion goes wrong."""


class RegistrationInfo(object):
    
    """Stored registration information about a model."""
    
    __slots__ = "fields", "file_fields", "follow", "format",
    
    def __init__(self, fields, file_fields, follow, format):
        """Initializes the registration info."""
        self.fields = fields
        self.file_fields = file_fields
        self.follow = follow
        self.format = format
        
        
class Registry(object):
    
    """A store of model registration information."""
    
    __slots__ = "_registry",
    
    def __init__(self):
        """Initializes the registry."""
        self._registry = {}
    
    def register(self, model_class, fields=None, follow=(), format="xml"):
        """Registers a model with this revision manager."""
        # Prevent multiple registration.
        if self.is_registered(model_class):
            raise RegistrationError, "%r has already been registered with Reversion." % model_class
        # Calculate serializable model fields.
        opts = model_class._meta
        local_fields = opts.local_fields + opts.local_many_to_many
        if fields is None:
            fields = [field.name for field in local_fields]
        fields = tuple(fields)
        # Calculate serializable model file fields.
        file_fields = []
        for field in local_fields:
            if isinstance(field, models.FileField) and field.name in fields:
                field.storage = VersionFileStorageWrapper(field.storage)
                file_fields.append(field)
        file_fields = tuple(file_fields)
        # Register the generated registration information.
        follow = tuple(follow)
        registration_info = RegistrationInfo(fields, file_fields, follow, format)
        self._registry[model_class] = registration_info
    
    def is_registered(self, model_class):
        """
        Checks whether the given model has been registered with this revision
        manager.
        """
        return model_class in self._registry
    
    def assert_registered(self, model_class):
        """Throws an exception if the given model class is not registered."""
        if not self.is_registered(model_class):
            raise RegistrationError, "%r has not been registered with Reversion." % model_class
    
    def get_registration_info(self, model_class):
        """Returns the registration information for the given model class."""
        self.assert_registered(model_class)
        return self._registry[model_class]
        
    def get_revision_set(self, object_set):
        """
        Follows all the registered relationships in the given set of objects to
        return a set containing the original objects plus all their related
        objects.
        """
        revision_set = set()
        def follow_relationships(obj):
            # Prevent recursion.
            if obj in revision_set:
                return
            revision_set.add(obj)
            # Follow relations.
            registration_info = self.get_registration_info(obj.__class__)
            for relationship in registration_info.follow:
                # Clear foreign key cache.
                try:
                    related_field = obj._meta.get_field(relationship)
                except models.FieldDoesNotExist:
                    pass
                else:
                    if isinstance(related_field, models.ForeignKey):
                        if hasattr(obj, related_field.get_cache_name()):
                            delattr(obj, related_field.get_cache_name())
                # Get the referenced object(s).
                related = getattr(obj, relationship, None)
                if isinstance(related, models.Model):
                    follow_relationships(related) 
                elif isinstance(related, (models.Manager, QuerySet)):
                    for related_obj in related.all():
                        follow_relationships(related_obj)
                elif related is not None:
                    raise TypeError, "Cannot follow the relationship %r. Expected a model or QuerySet, found %r." % (relationship, related)
        map(follow_relationships, object_set)
        return revision_set
        
    def unregister(self, model_class):
        """Removes a model from version control."""
        self.assert_registered(model_class)
        registration_info = self._registry.pop(model_class)
        for field in registration_info.file_fields:
            field.storage = field.storage.wrapped_storage
        
        
# A shared store of model registration information.
registry = Registry()

